import os
import json
import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
from utils.model import HorizonNet
from utils.dataset import visualize_a_data
from utils.misc import post_proc, utils
from utils.inference import inference


def run_inference(image_path, output_dir, model_path, flip=False, rotate=[], visualize=False,
                  force_cuboid=False, force_raw=False, min_v=None, r=0.05, device='cuda'):
    """
    封装推理脚本为函数，对单张图片进行推理并输出结果到指定目录。

    :param image_path: 要处理的图片路径
    :param output_dir: 输出目录路径
    :param model_path: 模型权重文件路径
    :param flip: 是否进行左右翻转增强，默认为 False
    :param rotate: 水平旋转增强参数，默认为空列表
    :param visualize: 是否可视化结果，默认为 False
    :param force_cuboid: 是否强制生成立方体布局，默认为 False
    :param force_raw: 是否强制输出原始多边形，默认为 False
    :param min_v: 峰值检测的最小值，默认为 None
    :param r: 峰值检测的半径比例，默认为 0.05
    :param device: 设备类型，默认为 'cuda'
    :return: 无返回值，结果直接保存到输出目录
    """
    # 检查输入图片是否存在
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"图片文件 {image_path} 不存在")

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 加载模型
    device = torch.device(device)
    net = utils.load_trained_model(HorizonNet, model_path).to(device)
    net.eval()

    # 加载图片
    img_pil = Image.open(image_path)
    if img_pil.size != (1024, 512):
        img_pil = img_pil.resize((1024, 512), Image.BICUBIC)
    img_ori = np.array(img_pil)[..., :3].transpose([2, 0, 1]).copy()
    x = torch.FloatTensor([img_ori / 255])

    # 推理
    with torch.no_grad():
        cor_id, z0, z1, vis_out = inference(net=net, x=x, device=device,
                                            flip=flip, rotate=rotate,
                                            visualize=visualize,
                                            force_cuboid=force_cuboid,
                                            force_raw=force_raw,
                                            min_v=min_v, r=r)

    # 保存结果
    basename = os.path.splitext(os.path.basename(image_path))[0]
    with open(os.path.join(output_dir, f"{basename}.json"), 'w') as f:
        json.dump({
            'z0': float(z0),
            'z1': float(z1),
            'uv': [[float(u), float(v)] for u, v in cor_id],
        }, f)

    # 保存可视化结果
    if vis_out is not None:
        vis_path = os.path.join(output_dir, f"{basename}.raw.png")
        vh, vw = vis_out.shape[:2]
        Image.fromarray(vis_out).resize((vw // 2, vh // 2), Image.LANCZOS).save(vis_path)

    print(f"推理完成，结果保存在: {output_dir}")
